import torch
import onnxruntime

# TODO: 根据最终模型结构修改此函数
def pt2onnx(model_path, onnx_path):
    model = torch.load(model_path, weights_only=False)
    model.eval()

    # 创建 Transformer 所需的 dummy_input
    dummy_src = torch.randint(0, 4, (1, 50), dtype=torch.long)  # (batch_size, seq_len)
    dummy_tgt = torch.randint(0, 4, (1, 50), dtype=torch.long)  # (batch_size, seq_len)
    dummy_src_mask = torch.ones(1, 1, 50, 50, dtype=torch.bool)
    dummy_tgt_mask = torch.ones(1, 1, 50, 50, dtype=torch.bool)

    torch.onnx.export(
        model,
        (dummy_src, dummy_tgt, dummy_src_mask, dummy_tgt_mask),
        onnx_path,
        opset_version=11
    )
    print("ONNX model saved to:", onnx_path)
    onnx_model = onnxruntime.InferenceSession(onnx_path)
    print("ONNX model loaded successfully.")


if __name__ == "__main__":
    model_path = r"D:\FengGong_Project\Machine_Trans\checkpoints\epoch_1.pt"
    onnx_path = r"D:\FengGong_Project\Machine_Trans\checkpoints\model.onnx"
    pt2onnx(model_path, onnx_path)
